#include "gtest/gtest.h"
#include "sgPetsc.h"
#include "Space.h"

//#include "DataStructure/Common/include/Constant.h"

//using SG::DataStructure::Common::TOLERANCE;
using namespace SG::Algebra;

class PetscTest : public ::testing::Test
{
public:
    static void SetUpTestSuite ()
    {

    }

    static void TearDownTestSuite() 
    {
    }
};

/** Test for parallel vector */
TEST_F (PetscTest, VectorTest)
{
    PetscVector<Real_t> v(10);
    v.assemble ();

#ifdef SG_VERBOSE
    std::cout << v << std::endl;
#endif
}

TEST_F (PetscTest, create_subvector_GeneralIndexSet_test)
{
    Index_t             L = 9;
    Real_t              coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscVector<Real_t> v(L);
    for (int i = 0; i < L; ++i)
    {
        // int I = local_to_global (i);
        int I = i;
        v.set (I, coeffs[i]);
    }
    v.assemble ();
    GeneralIndexSet v_set = {0, 1};
    PetscVector<Real_t> out_v;
    v.create_subvector(out_v, v_set);
// #ifdef SG_VERBOSE
    std::cout << out_v << std::endl;
// #endif
}

TEST_F (PetscTest, create_subvector_StrideIndexSet_test)
{
    Index_t             L        = 9;
    Real_t              coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscVector<Real_t> v (L);
    for (int i = 0; i < L; ++i)
    {
        // int I = local_to_global (i);
        int I = i;
        v.set (I, coeffs[i]);
    }
    v.assemble ();
    StrideIndexSet v_stride (2, 0, 1);
    
    PetscVector<Real_t> out_v;
    v.create_subvector (out_v, v_stride);
    // #ifdef SG_VERBOSE
    std::cout << out_v << std::endl;
    // #endif
}

TEST_F (PetscTest, linfty_test)
{
    Index_t             L        = 9;
    Real_t              coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscVector<Real_t> v (L);
    for (int i = 0; i < L; ++i)
    {
        // int I = local_to_global (i);
        int I = i;
        v.set (I, coeffs[i]);
    }
    v.assemble ();
    Real_t x;
    v.linfty_norm(x);
    // #ifdef SG_VERBOSE
    // std::cout << v << std::endl;
    std::cout << x << std::endl;
    // #endif
}

/** Test for parallel full matrix */
TEST_F (PetscTest, DenseMatrixTest)
{
    PetscDenseMatrix<Real_t> A (10, 10);

#ifdef SG_VERBOSE
    std::cout << A << std::endl;
#endif
}

/** Test for parallel sparse matrix */
TEST_F (PetscTest, SparseMatrixTest)
{
    PetscSparseMatrix<Real_t> A (10,10);
    A.assemble ();

#ifdef SG_VERBOSE
    std::cout << A << std::endl;
#endif
}

TEST_F (PetscTest, SubmatrixTest)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    StrideIndexSet rows_stride (1, 0, 1);  // n=0, first=0, step=1
    StrideIndexSet cols_stride (1, 0, 1);  // n=0, first=0, step=1 

    PetscSparseMatrix<Real_t> submatrix;
    mat.create_submatrix_nosort (submatrix, rows_stride, cols_stride);

#ifdef SG_VERBOSE
    std::cout << submatrix << std::endl;
#endif
}

TEST_F (PetscTest, Submatrix_GeneralIndexSet_Test)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    // std::vector<Index_t> rows_vet(n);
    GeneralIndexSet           rows = { 0, 1 };
    GeneralIndexSet           cols = { 0, 1 };
    PetscSparseMatrix<Real_t> submatrix;
    mat.create_submatrix_nosort (submatrix, rows, cols);

#ifdef SG_VERBOSE
    std::cout << submatrix << std::endl;
#endif
}

TEST_F (PetscTest, diagonal_Test)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 6.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    PetscVector<Real_t> dia_vec(M);
    mat.get_diagonal(dia_vec);
// #ifdef SG_VERBOSE
    std::cout << mat << std::endl;
    std::cout << dia_vec<< std::endl;
// #endif
}

TEST_F (PetscTest, inverse_Test)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 6.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    PetscSparseMatrix<Real_t> inverse_matrix;
    mat.get_inverse(inverse_matrix);
    // #ifdef SG_VERBOSE
    std::cout << mat << std::endl;
    std::cout << inverse_matrix << std::endl;
    // #endif
}

TEST_F (PetscTest, TransposeTest)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t  coeffs[] = { 6.0, 7.0, 5.0, 6.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    PetscSparseMatrix<Real_t> matT;
    mat.get_transpose (matT);

#ifdef SG_VERBOSE
    std::cout << mat << std::endl;
    std::cout << matT << std::endl;
#endif
}

TEST_F (PetscTest, scale_Test)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 6.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    const double s = 2.0;
    mat.scale(s);
// #ifdef SG_VERBOSE
    std::cout << mat << std::endl;
// #endif
}

TEST_F (PetscTest, mm_mult_Test)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count
    // Index_t m = 2;
    // Index_t n = 2;

    Real_t                    coeffs[] = { 6.0, 7.0, 5.0, 6.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    PetscSparseMatrix<Real_t> mat (M, N);

    auto row_start = mat.row_start ();
    auto row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();
    PetscSparseMatrix<Real_t> out_mat;
    mat.matrix_matrix_mult(mat, out_mat);
    // #ifdef SG_VERBOSE
    std::cout << out_mat << std::endl;
    // #endif
}
    /*
     * Test the A*x = b given by SuperLU User Guide.
     * where A is a 5X5 matrix and b is a 5X1 matrix.
     *
     * A =
     *   [
     *       06.00,  07.00,  05.00;
     *       07.00,  13.00,  08.00;
     *       05.00,  08.00,  06.00
     *   ]
     *
     * b = [9.0; 10.0; 9.0]
     *
     * OctaveOnline give the solution as followings,
     * x = inv(A)*b = [1; -1; 2]
     *
     * see https://octave-online.net/
     */
    TEST_F (PetscTest, PetscKspTest)
    {
        Index_t M = 3;  // row count
        Index_t N = 3;  // column count

        Real_t coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
        Real_t rhs[]    = { 9.0, 10.0, 9.0 };
        Real_t exact[]  = { 1, -1, 2 };

        PetscSparseMatrix<Real_t> mat (M, N);
        auto                      row_start = mat.row_start ();
        auto                      row_end   = mat.row_end ();

        for (auto i = row_start; i < row_end; ++i)
        {
            for (auto j = 0; j < N; ++j)
            {
                mat.set (i, j, coeffs[3 * i + j]);
            }
        }
        mat.assemble ();

        PetscVector<Real_t> b (M);
        for (auto i = row_start; i < row_end; ++i)
        {
            b.set (i, rhs[i]);
        }
        b.assemble ();

        TPetscKsp<Real_t> solver;
        solver.initialize ();
        solver.compute (mat);

        // Check
        auto info1 = solver.info ();
        ASSERT_TRUE (Info_t::Success == info1);

        // solution from OctaveOnline
        // x = inv(A)*b
        // see https://octave-online.net/
        PetscVector<Real_t> x (M);
        solver.solve (b, x);

#ifdef SG_VERBOSE
    std::cout << mat << std::endl;
    std::cout << b << std::endl;
    std::cout << x << std::endl;
#endif

    // Check
    auto info2 = solver.info ();
    ASSERT_TRUE (Info_t::Success == info2);

    for (auto i = row_start; i < row_end; ++i)
    {
        ASSERT_NEAR (exact[i], x (i), TOLERANCE);
    }
}

TEST_F (PetscTest, PetscKspComplexTest)
{
    Index_t M = 3;  // row count
    Index_t N = 3;  // column count

    Complex_t coeffs[] = { 6.0, 7.0, 5.0, 7.0, 13.0, 8.0, 5.0, 8.0, 6.0 };
    Complex_t rhs[]    = { 9.0, 10.0, 9.0 };
    Complex_t exact[]  = { 1, -1, 2 };

    PetscSparseMatrix<Complex_t> mat (M, N);
    auto           row_start = mat.row_start ();
    auto           row_end   = mat.row_end ();

    for (auto i = row_start; i < row_end; ++i)
    {
        for (auto j = 0; j < N; ++j)
        {
            mat.set (i, j, coeffs[3 * i + j]);
        }
    }
    mat.assemble ();

    PetscVector<Complex_t> b (M);
    for (auto i = row_start; i < row_end; ++i)
    {
        b.set (i, rhs[i]);
    }
    b.assemble ();

    TPetscKsp<Complex_t> solver;
    solver.initialize ();
    solver.compute (mat);

    // Check
    auto info1 = solver.info ();
    ASSERT_TRUE (Info_t::Success == info1);

    // solution from OctaveOnline
    // x = inv(A)*b
    // see https://octave-online.net/
    PetscVector<Complex_t> x (M);
    solver.solve (b, x);

#ifdef SG_VERBOSE
    std::cout << mat << std::endl;
    std::cout << b << std::endl;
    std::cout << x << std::endl;
#endif

    // Check
    auto info2 = solver.info ();
    ASSERT_TRUE (Info_t::Success == info2);

    for (auto i = row_start; i < row_end; ++i)
    {
        ASSERT_NEAR (exact[i].real (), x (i).real (), TOLERANCE);
        ASSERT_NEAR (exact[i].imag (), x (i).imag (), TOLERANCE);
    }
}
